package middleware

import (
	"context"
	"io"
	"net/http"

	"golang.org/x/time/rate"
)

type UploadRateLimitMiddleware struct {
	bytesPerSecond int
}

func NewUploadRateLimitMiddleware(bytesPerSecond int) *UploadRateLimitMiddleware {
	return &UploadRateLimitMiddleware{
		bytesPerSecond: bytesPerSecond,
	}
}

func (m *UploadRateLimitMiddleware) Handle(next http.HandlerFunc) http.HandlerFunc {
	limiter := rate.NewLimiter(
		rate.Limit(m.bytesPerSecond),
		m.bytesPerSecond*2, // 突发限制为2倍速率
	)

	return func(w http.ResponseWriter, r *http.Request) {
		reader := &RateLimitedReader{
			r.Context(),
			r.Body,
			limiter,
		}
		r.Body = reader

		// Passthrough to next handler if need
		next(w, r)
	}
}

type RateLimitedReader struct {
	ctx     context.Context
	src     io.ReadCloser
	limiter *rate.Limiter
}

func (r *RateLimitedReader) Read(p []byte) (n int, err error) {
	n, err = r.src.Read(p)
	if err := r.limiter.WaitN(r.ctx, n); err != nil {
		return 0, err
	}
	return n, err
}

func (r *RateLimitedReader) Close() error {
	return r.src.Close()
}
